import logging
from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from tqdm import tqdm

# 设置日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 初始化Pinecone客户端
pinecone = Pinecone(api_key="0300a7a1-91e1-4ddb-99c9-15c631060bb2")

# 索引名称
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.info(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 分割数据集为训练集和测试集
split_index = int(0.8 * len(X))
train_X, test_X = X[:split_index], X[split_index:]
train_y, test_y = y[:split_index], y[split_index:]

# 初始化向量列表
vectors = []

# 处理训练数据
for i, (x, label) in enumerate(zip(train_X, train_y)):
    vector_id = str(i)
    vector_values = x.tolist()
    metadata = {"label": int(label)}
    vectors.append((vector_id, vector_values, metadata))

# 上传数据到Pinecone
batch_size = 1000
total_vectors = len(vectors)
for i in tqdm(range(0, total_vectors, batch_size), desc="上传数据到Pinecone", total=(total_vectors + batch_size - 1) // batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
logging.info(f"成功上传了 {total_vectors} 条数据到 Pinecone")

# 测试准确率
correct_count = 0
total_count = len(test_X)

for i, (x, label) in tqdm(enumerate(zip(test_X, test_y)), total=total_count, desc="测试k=11的准确率"):
    query_data = x.tolist()
    results = index.query(vector=query_data, top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]
    
    if labels:  # 检查labels列表是否为空
        final_prediction = Counter(labels).most_common(1)[0][0]
        if final_prediction == label:
            correct_count += 1

accuracy = correct_count / total_count
logging.info(f"当k=11时，使用Pinecone的准确率为：{accuracy*100:.2f}%")

# 显示一个查询图像和预测结果
query_image = test_X[0]
query_label = test_y[0]
query_data = query_image.tolist()
results = index.query(vector=query_data, top_k=11, include_metadata=True)
labels = [match['metadata']['label'] for match in results['matches']]

if labels:  # 检查labels列表是否为空
    final_prediction = Counter(labels).most_common(1)[0][0]
else:
    final_prediction = None
    logging.info("无法预测结果，labels列表为空。")

# 使用 matplotlib 显示查询图像和预测结果
plt.imshow(query_image.reshape(8, 8), cmap='gray')
plt.title(f"Prediction: {final_prediction}, True Label: {query_label}")
plt.show()
